(*
file : inference.ml
Arnaud Aujon   AUJA31018906
*)

open Variable
open Type
open List

exception TypeError;;

let in_if = ref false;;

let rec rename_var old_name new_name expr =
  match expr with
    T2_Var(name) ->  (if (String.compare name old_name == 0)
          then begin print_string (name^" > "^new_name^"\n"); (T2_Var(new_name))
          end else expr)
  | T2_Binary(e1, op, e2) -> let new_e1 = rename_var old_name new_name e1
              and new_e2 = rename_var old_name new_name e2 in
              (T2_Binary(new_e1, op, new_e2))
  | T2_Unary(op, e) -> (T2_Unary(op, rename_var old_name new_name e))
  | T2_Assignment(var, e) -> let new_e = rename_var old_name new_name e in
              (T2_Assignment(rename_var old_name new_name var, new_e))
  | T2_Struct_If(if_list) ->
      (T2_Struct_If(List.map (rename_var_if old_name new_name) if_list ))
  | T2_Print(var_list) -> T2_Print(List.map (rename_var old_name new_name ) var_list)
  | T2_Return(e) -> T2_Return(rename_var old_name new_name e)
  | T2_Function(name, args) ->
      T2_Function(name, List.map (rename_var old_name new_name) args)
  | T2_For(var, b, e, body) ->
    let new_var = rename_var old_name new_name var
    and new_b = rename_var old_name new_name b
    and new_e = rename_var old_name new_name e
    and new_body = List.map (rename_var old_name new_name) body in
    T2_For(new_var, new_b, new_e, new_body)
  | T2_While(cond, body) ->
    let new_cond = rename_var old_name new_name cond
    and new_body = List.map (rename_var old_name new_name) body in
    T2_While(new_cond, new_body)
  | T2_Decl_Function(name, args, body) -> T2_Decl_Function(name, args, body)
  | _ -> expr
  

and rename_var_if old_name new_name expr =
  match expr with
    T2_If(cond, expr_list) ->
      let new_cond = rename_var old_name new_name cond and
      new_expr = (List.map (rename_var old_name new_name) expr_list) in
        (T2_If(new_cond, new_expr))
  | T2_Else(expr_list) ->
      let new_expr = List.map (rename_var old_name new_name) expr_list in
        (T2_Else(new_expr))
;;

let rec get_type_string t =
  match t with
    TInt -> "int"
  | TFloat -> "float"
  | TBool -> "bool"
  | TString -> "string"
  | TArray -> "array"
  | TNil -> "all"
  | TObjVar(t2) -> "obj"^(get_type_string t2)
  | TArray_index -> "todo" 
  | TArray_int -> error "unexpected type in get_type_string"
;;

(* convert t2_type to t3_type *)
let rec t2_to_t3 e =
  match e with
    T2_Var(name) -> T3_Var(name)
  | T2_Binary(e1, op, e2) -> T3_Binary(t2_to_t3 e1, op, t2_to_t3 e2)
  | T2_Unary(op, e) -> T3_Unary(op, t2_to_t3 e)
  | T2_Int(v) -> T3_Int(v)
  | T2_Float(v) -> T3_Float(v)
  | T2_Bool(v) -> T3_Bool(v)
  | T2_String(v) -> T3_String(v)
  | T2_Array(ls) -> T3_Array(List.map t2_to_t3 ls)
  | T2_Array_index(_,_) -> error "this is an limitation"
  | T2_Print(var_list) -> T3_Print(List.map t2_to_t3 var_list)
  | T2_Return(e) -> T3_Return(t2_to_t3 e)
  | T2_Function(name, args) -> T3_Function(name, List.map t2_to_t3 args)
  | T2_For(var, b, e, body) ->
    let new_var = t2_to_t3 var
    and new_b = t2_to_t3 b
    and new_e = t2_to_t3 e
    and new_body = List.map t2_to_t3 body in
    T3_For(new_var, new_b, new_e, new_body)
  | T2_While(cond, body) ->
    let new_cond = t2_to_t3 cond
    and new_body = List.map t2_to_t3 body in
    T3_While(new_cond, new_body)
  | _ -> error "unmatched structure in t2_to_t3 ()"
;;

(* return true is expr is an objVar *)
let rec expr_is_objVar expr =
  match expr with
    T2_Function(name, args) -> expr_is_objVar (T2_Var(name))
  | T2_Var(_) ->
    begin match get_var_type (get_name_t2 expr) with
      TObjVar(_) -> true
    | TNil -> true
    | _ -> false
    end
  | _ -> false
;;


(* return type of an expression *)
let rec get_expr_type expr =
  match expr with
      T2_Int(_) -> TInt
    | T2_Float(_) -> TFloat
    | T2_Bool(_) -> TBool
    | T2_String(_) -> TString
    | T2_Array(_) -> TArray
    | T2_Binary(e1,op,e2) ->
    let t1 = get_expr_type e1
    and t2 = get_expr_type e2 in
    let t =
      if (t_is_obj t1) then t2
      else if (t_is_obj t2) then t1
      else if (t1 = TNil && t2 != TNil) then t2
      else if (t1 != TNil && t2 = TNil) then t2
      else TObjVar(TNil) in

      begin match op with
        "&&" -> if ((t1 = TBool || t_is_obj t1) && (t2 = TBool || t_is_obj t2))
                then TBool else raise TypeError
      | "||" -> if ((t1 = TBool || t_is_obj t1) && (t2 = TBool || t_is_obj t2))
                then TBool else raise TypeError
      | "==" -> t
      | "!=" -> t
      | "<" -> t
      | ">" -> t
      | "<=" -> t
      | ">=" -> t
      | "+" -> t
      | "-" -> t
      | "*" -> t
      | "**" -> if (t_is_obj t1) then TObjVar(t2) else TObjVar(t1)
      | "/" -> t
      | _ -> error ("undefined binary operator `"^op^"`\n")
      end
    | T2_Function(name, arg_list) -> get_var_type name
    | T2_Var(_) ->  get_var_type (get_name_t2 expr)
  | _ -> error "error in get_expr_type : wird expr"
;;


let rec type_expr t expr =
  match expr with
    T2_Int(v) -> ()
  | T2_Float(v) -> ()
  | T2_Bool(v) -> ()
  | T2_String(v) -> ()
  | T2_Array(ls) -> ()
  | T2_Array_index(name,idx) -> ()
  | T2_Print(var_list) -> ()
  | T2_Var(name) -> print_string ("in type_expr :"^name^"\n");
                    let t2 = get_expr_type expr in
                    if ( t2 = TNil ) then begin
						print_string "change_type !\n";
						rm_var name;
					    add_var t name;
					end;
					()
  | T2_Binary(e1, op, e2) -> type_expr t e1; type_expr t e2
  | T2_Unary(op, e) -> type_expr t e
  | _ -> ()
;;

(* perform type inference which will surround all variable assignment
by an indication of there type *)
let rec inference_type ast =
  match ast with
      [] -> []
    | head::tail -> 
      begin match head with
        T2_Assignment(var, e) ->
        let name = get_name_t2 var in
          print_string ("assign `"^name^"`\n");
          let t_var = get_var_type name and t_expr = get_expr_type e in
          print_string ("type : "^get_type_string t_expr);
          type_expr t_expr e;
          let new_e = e in
          if ( (t_var = t_expr) || (t_var = TNil) )
          then begin (* declaration is same type *)
            print_string "same_type\n";
            rm_var name;
            begin match t_expr with
              TObjVar(_) -> print_string "obj\n"; obj := name::!obj ; add_var t_expr name
           
            | t2 -> add_var t2 name
            end;
			let new_v = t2_to_t3 var in
			let new_assgn=T3_Typed_Assignment(t_expr, new_v, t2_to_t3 new_e) in
            new_assgn::(inference_type tail)
          end else begin      (* different type -> it's an object ! *)
            if ( !in_if ) (* we are in a if *)
              then begin
              print_string "in if\n";
              rm_var name;
              add_var (TObjVar(t_expr)) name;
              obj := name::!obj;
              T3_Typed_Assignment(t_expr,
                t2_to_t3 var, t2_to_t3 new_e)
              ::(inference_type tail)
            end else begin
              if ( t_is_obj t_expr )
                then begin
                  print_string "transofrm_to_obj\n";
                  rm_var name;
                  add_var t_expr name;
                  obj := name::!obj;
                  T3_Typed_Assignment(t_expr,
                    t2_to_t3 var, t2_to_t3 new_e)
                  ::(inference_type tail)
              end else begin
                (* OPTIMISATION POSSIBLE ICI *)
                print_string ("opt !"^name^"\n");
                (* if name is an object remove from obj list *)
                let rec rm_obj = function [] -> []
                             | head::tail -> if (String.compare head name == 0)
                                               then tail
                                             else head::(rm_obj tail) in
                obj := rm_obj !obj;

                let new_name = get_new_name () in
                add_var t_expr new_name;
                let new_tail = List.map (rename_var name new_name) tail in
                T3_Typed_Decl_Var(new_name, t_expr)::
                T3_Typed_Assignment(t_expr, T3_Var(new_name),t2_to_t3 new_e)
                  ::(inference_type new_tail)
              end
            end
          end
      | T2_Struct_If(if_list) ->
        (*
        la variable in_if est utilisée pour spécifier que le type type inféré
        dans un if n'est pas forcement le type "unique" de la variable
        *)
        in_if := true;
        let new_if_list = List.map inference_if if_list in
        in_if := false;
        T3_Struct_If(new_if_list)::(inference_type tail)

      | T2_For(var, beg, end_, body) ->
        let t = get_expr_type beg in
        let name = get_name_t2 var in
        rm_var name;
        add_var t name;
        let new_body = List.map inference_type [body] in
        T3_For(t2_to_t3 var, t2_to_t3 beg, t2_to_t3 end_, hd new_body)
          ::(inference_type tail)
      | T2_While(cond, body) -> 
        (*let new_cond = *)
        let new_body = List.map inference_type [body] in
        T3_While(t2_to_t3 cond, hd new_body)::(inference_type tail)
      | T2_Decl_Var(name) -> add_var TNil name ; T3_Decl_Var(name)
          ::(inference_type tail)
      | T2_Decl_Function(name, args, body) ->
       (* print_string ("decl_fun : "^name);*)
        let old_env = !env in
          env := [];
          List.iter (function x -> (add_var TNil x)) args;
          let new_body = inference_type body in
          let t_r = get_var_type "Return" in
            let f = function x ->
                let t = get_var_type x in
                print_string ("arg : "^x^" : "^(get_type_string t));       
                if ( t = TNil ) then (x, TObjVar(TNil))
                else (x, t)
            in
          let new_args = List.map f args in
          let f2 = function
            T3_Decl_Var(n) -> let t = get_var_type n in
                                (T3_Typed_Decl_Var(n, t)) 
            | x -> x in
          let new_body2 = List.map f2 new_body in
        env := old_env;
        (* pb : si le type de retour ne peut être inféré,
        on a un TNil alors qu'il faut un ObjVar(TNil) 
        TODO : find a better fix *)
        let t_r2 = if ( t_r = TNil ) then TObjVar(TNil) else t_r in
        
        add_var t_r2 name;
        print_string ((get_type_string t_r2)^"\n");
        
        (T3_Typed_Decl_Function(name, new_args , new_body2)
          ::(inference_type tail))
      | T2_Return(e) ->
        (*print_string "return\n";*)
        let t = get_expr_type e in add_var t "Return";
        T3_Return(t2_to_t3 e)::(inference_type tail)
      | T2_Int(v) -> T3_Int(v)::(inference_type tail)
      | T2_Float(v) -> T3_Float(v)::(inference_type tail)
      | T2_Bool(v) -> T3_Bool(v)::(inference_type tail)
      | T2_String(v) -> T3_String(v)::(inference_type tail)
      | T2_Function(name, args) ->
        T3_Function(name, List.map t2_to_t3 args)::(inference_type tail)
      | x -> (t2_to_t3 x)::(inference_type tail)
      end

and inference_if if_list =
  match if_list with
    T2_If(cond, body) -> let new_body = inference_type body in  
              T3_If(t2_to_t3 cond,new_body)
  | T2_Else(body) -> let new_body = inference_type body in 
            T3_Else(new_body)
;;